import pygame
import ale_py
import gymnasium as gym
import time

# 初始化环境和手柄
gym.register_envs(ale_py)
env = gym.make("ALE/Breakout-v5", render_mode="human")
env.reset()

pygame.init()
pygame.joystick.init()
joystick = pygame.joystick.Joystick(0)
joystick.init()

# 动作映射（根据手柄实际按键调整）
ACTION_MAP = {
    "left_axis_x_negative": 3,  # 左移
    "left_axis_x_positive": 2,  # 右移
    "button_0": 1               # 发射（通常为A键）
}

running = True
action = 0  # 默认无动作

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
    action = 0  # 重置动作
    # 检测摇杆输入（左摇杆X轴）
    axis_x = joystick.get_axis(0)  # 通常左摇杆X轴为轴0
    if axis_x < -0.5:
        action = ACTION_MAP["left_axis_x_negative"]
    elif axis_x > 0.5:
        action = ACTION_MAP["left_axis_x_positive"]

    if joystick.get_hat(0) == (1, 0):  # 检测右方向键
        action = ACTION_MAP["left_axis_x_positive"]
    elif joystick.get_hat(0) == (-1, 0):  # 检测左方向键
        action = ACTION_MAP["left_axis_x_negative"]
    # 检测按钮输入（例如A键发射）
    if joystick.get_button(0):  # 通常按钮0对应A键
        action = ACTION_MAP["button_0"] 

    # 执行动作
    obs, reward, terminated, truncated, info = env.step(action)
    env.render()

    if terminated or truncated:
        print('游戏结束，按"返回"键退出...')
        time.sleep(0.1)
        # env.reset()

    # 按手柄的"返回"键退出（通常为按钮6）
    if joystick.get_button(6):
        running = False

env.close()